import numpy as np
from gem5.simulate import exit_event, simulate
from gem5.utils import overrides
from gem5.components.boards.simple_board import SimpleBoard
from gem5.components.memory.single_channel import SingleChannelDDR3_1600
from gem5.components.processors.simple_processor import SimpleProcessor
from gem5.components.processors.cpu_types import CPUTypes
from gem5.isas import ISA

class GasPipelineNetwork:
    def __init__(self, nodes, edges):
        self.nodes = nodes  # 节点列表（阀门/关键点）
        self.edges = edges  # 管道连接关系
        self.adj_matrix = self._build_adjacency_matrix()  # 邻接矩阵
    
    def _build_adjacency_matrix(self):
        """构建邻接矩阵表示管网连接关系"""
        size = len(self.nodes)
        matrix = np.zeros((size, size), dtype=int)
        for (i, j) in self.edges:
            matrix[i][j] = 1
            matrix[j][i] = 1  # 双向管道
        return matrix
    
    def warshall_transitive_closure(self):
        """Warshall算法计算传递闭包，识别管网子区"""
        closure = self.adj_matrix.copy()
        n = closure.shape[0]
        for k in range(n):
            for i in range(n):
                for j in range(n):
                    closure[i][j] = closure[i][j] or (closure[i][k] and closure[k][j])
        return closure
    
    def find_critical_valves(self, accident_node):
        """基于传递闭包确定需关闭的阀门"""
        closure = self.warshall_transitive_closure()
        affected_nodes = np.where(closure[accident_node] == 1)[0]
        # 找出隔离事故区域的最小阀门集合（此处简化为直接关闭相邻节点）
        critical_valves = []
        for node in affected_nodes:
            if self.adj_matrix[accident_node][node] == 1:
                critical_valves.append(node)
        return critical_valves

# 示例管网配置
nodes = ["V1", "V2", "V3", "V4", "V5"]
edges = [(0,1), (1,2), (2,3), (3,4), (0,4)]
pipeline = GasPipelineNetwork(nodes, edges)